#!/usr/bin/env python3
import sys
import os
sys.path.append(os.path.join(os.path.dirname(__file__), '..', 'debugger'))
import numpy as np
from semPcd import semPcd
from scipy.spatial.transform import Rotation
from transPcd import Transform
from termcolor import cprint
from time_it import timeit
from scalePcd import scalePcd
from save_file import save_pcd
from termcolor import colored
import open3d as o3d
from loadJson import _load_finger_length, _load_emb_name_map
class ENGPRPCD:
    def __init__(self, unifiedData: dict, param: str, pcd_keys: str, surpress_print: bool = False, DEBUG_PCD: bool = True,
                 seg_source: bool = False, episode: int | None = None, ts: int | None = None):
        self.unifiedData = unifiedData
        assert param[0] in ['b','0'], f"enGprPcd: param {param} must start with 'b' or '0'"
        self.show_bbox = (param[0]=='b') if len(param)>=1 else False
        assert len(param)>=2, f"enGprPcd: param {param} must have at least 2 characters"
        self.embodiment_char = param[1]
        self.neutral_stretch_file = None
        self.transform_filename = None
        self.save_neutral_stretch = False
        self.save_local_gripper = False
        self.targ_embodiment_char = None
        self.contrastive_color = False
        remaining_param = param[2:] if len(param) > 2 else ""
        if remaining_param:
            # Remove leading underscore if present and split
            if remaining_param.startswith('_'):
                remaining_param = remaining_param[1:]
            parts = remaining_param.split('_')
            if len(parts) >= 1:
                self.save_neutral_stretch = (parts[0]=='1')
                self.neutral_stretch_file = parts[0] if parts[0] not in ["0","1"] else None
            if len(parts) >= 2:
                self.transform_filename = parts[1] if parts[1] not in ["0","1"] else None
            if len(parts) >= 3:
                self.save_local_gripper=(parts[2]=='1' and ts==0)
                self.targ_embodiment_char=parts[2] if parts[2] not in ["0","1"] else None
            if len(parts) >= 4:
                self.contrastive_color=(parts[3]=='c')
            
        print(f"save_neutral_stretch {self.save_neutral_stretch} neutral_stretch_file {self.neutral_stretch_file} transform_filename {self.transform_filename} save_local_gripper {self.save_local_gripper} targ_embodiment_char {self.targ_embodiment_char}")
        self.surpress_print = surpress_print
        self.DEBUG_PCD = DEBUG_PCD
        self.seg_source = seg_source
        self.pcd_keys=pcd_keys
        self.episode=episode
        self.ts=ts
        self.key_poses = unifiedData.get('key_poses', None)
        self.semantic_indices = semPcd(unifiedData, self.embodiment_char, surpress_print=True, seg_source=seg_source)
        # print(f"self.semantic_indices: {self.semantic_indices.keys()}")
        self.gripper_parts = ['lhand_lfinger','lhand_rfinger','lhand', 'rhand_lfinger', 'rhand_rfinger','rhand']
        self.gripper_indices = {part: self.semantic_indices[part] for part in self.gripper_parts if part in self.semantic_indices}
    def _get_bbox(self, local_points: np.ndarray) -> np.ndarray:
        x_min, y_min, z_min = local_points[:, 0].min(), local_points[:, 1].min(), local_points[:, 2].min()
        x_max, y_max, z_max = local_points[:, 0].max(), local_points[:, 1].max(), local_points[:, 2].max()
        return np.array([[x_min, y_min, z_min], [x_max, y_min, z_min], [x_min, y_max, z_min], [x_max, y_max, z_min],
                        [x_min, y_min, z_max], [x_max, y_min, z_max], [x_min, y_max, z_max], [x_max, y_max, z_max]])

    def _color_closest_4_points(self, bbox_8_points: np.ndarray, ref_point: np.ndarray) -> np.ndarray:
        bbox_8_points = np.column_stack([bbox_8_points, np.zeros((8, 3))]) if bbox_8_points.shape[1] == 3 else bbox_8_points
        distances = np.linalg.norm(bbox_8_points[:, :3] - ref_point[:3], axis=1)
        closest_indices = np.argsort(distances)[:4]
        bbox_8_points[closest_indices, 3:6] = [0, 255, 255]
        bbox_8_points[np.setdiff1d(np.arange(8), closest_indices), 3:6] = [255, 0, 0]
        return bbox_8_points

    def _keypose_center(self, part_name: str) -> np.ndarray:
        side = 'left' if part_name.startswith('lhand') else 'right' if part_name.startswith('rhand') else None
        assert side is not None, f"enGprPcd: Cannot determine side for part {part_name}"
        return self.key_poses[30:36] if side == 'left' else self.key_poses[12:18]
    def examine_4points_sameface(self, bbox_local_colored: np.ndarray) -> np.ndarray:
        cyan_mask = np.all(bbox_local_colored[:, 3:6] == [0, 255, 255], axis=1)
        cyan_points = bbox_local_colored[cyan_mask, :3]
        x_coords, y_coords, z_coords = cyan_points[:, 0], cyan_points[:, 1], cyan_points[:, 2]
        if len(np.unique(x_coords)) == 1: return np.array([1, 0, 0]) if x_coords[0] > 0 else np.array([-1, 0, 0])
        elif len(np.unique(y_coords)) == 1: return np.array([0, 1, 0]) if y_coords[0] > 0 else np.array([0, -1, 0])
        elif len(np.unique(z_coords)) == 1: return np.array([0, 0, 1]) if z_coords[0] > 0 else np.array([0, 0, -1])
        cprint(f"4 cyan points not on same face {cyan_points}",'red')
        return np.array([0, 0, 1])

    def _enhance_pcd(self, part_points_local: np.ndarray, norm_vecs: np.ndarray, bbox_local_colored: np.ndarray, part_name: str):
        enhanced_points = Transform.get_transforms(
            part_points_local,
            norm_vecs,
            bbox_local_colored,
            part_name,
            stick_faces=True,
            z_boundary=False,
            transform_file=self.transform_filename,
            neutral_stretch=self.neutral_stretch_file,
            want_neutral_stretch=self.save_neutral_stretch,
            episode=self.episode,
            ts=self.ts,
            emb_char=self.embodiment_char
        )
        if self.show_bbox:
            enhanced_points = np.vstack([enhanced_points, bbox_local_colored])
        enhanced_points_global = Transform.local_points_to_global(enhanced_points, self._keypose_center(part_name)[3:6],self._keypose_center(part_name)[:3])
        return enhanced_points_global
    def remove_gripper_points(self,reOrsegPcd:np.ndarray):
        all_gripper_indices = []
        for indices in self.gripper_indices.values():
            all_gripper_indices.extend(indices)
        all_gripper_indices = list(set(all_gripper_indices))
        mask = np.ones(reOrsegPcd.shape[0], dtype=bool)
        mask[all_gripper_indices] = False
        reOrsegPcd = reOrsegPcd[mask]
        return reOrsegPcd
    def get_parts_data(self) -> dict:
        parts_data = {}
        for part, indices in self.gripper_indices.items():
            assert indices, f"part {part} has no indices, indices={indices}"
            part_points = self.pcd[indices]
            part_points_local = Transform.global_points_to_local(part_points,self._keypose_center(part)[3:6],self._keypose_center(part)[:3])
            bbox_local = self._get_bbox(part_points_local)
            center_global = self._keypose_center(part)[:3]
            center_local = Transform.global_points_to_local(center_global.reshape(1, 3),self._keypose_center(part)[3:6],self._keypose_center(part)[:3])[0]
            assert np.allclose(center_local, [0,0,0]), f"center_local: {center_local} should be [0,0,0]"
            bbox_local_colored = self._color_closest_4_points(bbox_local, center_local)
            norm_vector = self.examine_4points_sameface(bbox_local_colored)
            dims = bbox_local.max(axis=0) - bbox_local.min(axis=0)
            dims = np.abs(dims)
            dims_sorted = np.sort(dims)[::-1]
            length, width, height = dims_sorted.tolist()
            # print(colored(f"{part} dims: {length:.4f} {width:.4f} {height:.4f}","cyan"))
            parts_data[part] = {'global':{'bbox': Transform.local_points_to_global(bbox_local_colored,self._keypose_center(part)[3:6],self._keypose_center(part)[:3]), 'center': center_global,'part_points': part_points},
                                'local' :{'bbox': bbox_local_colored,                                    'center': center_local, 'part_points': part_points_local,'norm_vector': norm_vector}}
            if not self.surpress_print: 
                print(f"    🔧 ENGRIP | {part}: {len(indices)} points → enhanced")
        return parts_data
    def enhance_parts(self,parts_data:dict):
        hand_norm={}
        enhanced_parts = []
        print(f"parts_data: {parts_data.keys()}")   
        for part, data in parts_data.items():
            if part.endswith('hand'):
                hand_norm[part[0]]=-data['local']['norm_vector']
        print(f"hand_norm: {hand_norm}")   
        for part, data in parts_data.items():
            norm_vecs=[data['local']['norm_vector']]
            if not part.endswith('hand') and part[0] in hand_norm:
                norm_vecs.append(hand_norm[part[0]])
            elif part[0] not in hand_norm:
                print(f"part {part[0]} from {part} not in hand_norm {hand_norm.keys()}")
            enhanced_parts.append(self._enhance_pcd(data['local']['part_points'], norm_vecs, data['local']['bbox'], part))
        return np.vstack(enhanced_parts) if enhanced_parts else np.array([])
    def check_re_seg_same_idx(self):    
        assert 'repointcloud' in self.unifiedData and 'segpointcloud' in self.unifiedData and self.unifiedData['repointcloud'] is not None and self.unifiedData['segpointcloud'] is not None, f"repointcloud or segpointcloud is None key:{self.unifiedData.keys()} re:{self.unifiedData.get('repointcloud',None)} seg:{self.unifiedData.get('segpointcloud',None)}"
        assert self.unifiedData['repointcloud'].shape[0] == self.unifiedData['segpointcloud'].shape[0], f"repointcloud ({self.unifiedData['repointcloud'].shape[0]} points) and segpointcloud ({self.unifiedData['segpointcloud'].shape[0]} points) must have the same number of points"
        assert np.allclose(self.unifiedData['repointcloud'][:, :3], self.unifiedData['segpointcloud'][:, :3], atol=1e-6), f"XYZ coordinates of points at same indices in repointcloud and segpointcloud must be identical"
    def add_pcd_to_result(self,pcd:np.ndarray,pcd_key:str):
        assert pcd_key in self.unifiedData
        self.unifiedData[pcd_key]=np.vstack([self.unifiedData[pcd_key],pcd])
        self.unifiedData[pcd_key]=self.remove_gripper_points(self.unifiedData[pcd_key])
        return self.unifiedData
    
    def save_local_gripper_pcd(self, parts_data: dict, key: str):
        left_gripper_parts = [data['local']['part_points'] for part_name, data in parts_data.items() if part_name.startswith('l')]
        assert left_gripper_parts, "No left gripper parts found to save"
        combined_local_points = np.vstack(left_gripper_parts)
        save_pcd(f"local_gripper_{key}_{self.embodiment_char}.pcd",combined_local_points)
    def compare_local_gripper_pcd(self, key: str) -> np.ndarray:
        pcd = o3d.io.read_point_cloud(f"local_gripper_{key}_{self.targ_embodiment_char}.pcd")
        local_points = np.asarray(pcd.points)
        colors = np.asarray(pcd.colors)
        pcd1_local=np.column_stack([local_points, colors])  
        pcd1_local=scalePcd(pcd1_local, tag="local_gripper", suppress_print=True)
        pcd1_global= Transform.local_points_to_global(pcd1_local, self._keypose_center('lhand')[3:6], self._keypose_center('lhand')[:3])
        pcd2_local=np.column_stack([local_points, colors])  
        pcd2_local=scalePcd(pcd2_local, tag="local_gripper", suppress_print=True)
        pcd2_global= Transform.local_points_to_global(pcd2_local, self._keypose_center('rhand')[3:6], self._keypose_center('rhand')[:3])
        return np.vstack([pcd1_global,pcd2_global])
    def enGprPcd(self) -> dict:
        self.check_re_seg_same_idx()
        for key in self.pcd_keys:
            assert key in ["repointcloud","segpointcloud"], f"enGprPcd: {key} is not a valid pcd key.currently only support repointcloud and segpointcloud"
            self.pcd=self.unifiedData[key]
            print(key,"maxcolor",self.pcd[:,3:6].max(axis=0))
            assert self.pcd.shape[0] > 0, f"enGprPcd: {key} is empty."
            if not self.gripper_indices:
                g_transformed_pcd = np.zeros((0,6))
                if self.save_local_gripper:
                    print(colored("enGprPcd: save_local_gripper is True, but no gripper parts found.","red"))
            else:
                parts_data = self.get_parts_data()
                g_transformed_pcd=self.enhance_parts(parts_data)
                self.unifiedData['stick_bias_vec']=Transform.stick_bias_vec
                if self.save_local_gripper:
                    self.save_local_gripper_pcd(parts_data,key)
            if self.targ_embodiment_char:
                g_compare_pcd=self.compare_local_gripper_pcd(key)
                emb_map=_load_emb_name_map()
                finger_info=_load_finger_length()
                cur_emb_color=np.array(finger_info[emb_map[self.embodiment_char]]['contrastive_color'])
                targ_emb_color=np.array(finger_info[emb_map[self.targ_embodiment_char]]['contrastive_color'])
                if self.contrastive_color:
                    g_compare_pcd[:,3:6]=targ_emb_color
                    g_transformed_pcd[:,3:6]=cur_emb_color
                g_transformed_pcd=np.vstack([g_transformed_pcd,g_compare_pcd])
            self.add_pcd_to_result(g_transformed_pcd,key)
        return self.unifiedData
@timeit()
def enGprPcd(unifiedData: dict, param: str, surpress_print: bool = False, DEBUG_PCD: bool = True, seg_source: str = "traverse",
             episode: int | None = None, ts: int | None = None) -> dict:
    engprpcd = ENGPRPCD(
        unifiedData,
        param,
        ['repointcloud', 'segpointcloud'],
        surpress_print=surpress_print,
        DEBUG_PCD=DEBUG_PCD,
        seg_source=seg_source,
        episode=episode,
        ts=ts,
    )
    return engprpcd.enGprPcd()